package com.lee.windowFunction;

import com.lee.entity.SessionResult;
import org.apache.flink.streaming.api.functions.windowing.AllWindowFunction;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.HashMap;

/**
 * Created by liruopeng on 2018/1/18.
 */
public class SessionStatisticalCountWindowFunction implements AllWindowFunction<SessionResult, SessionResult, TimeWindow> {


    private Logger logger = LoggerFactory.getLogger(SessionStatisticalCountWindowFunction.class);

    private HashMap<String, Integer> countDstMap;
    private HashMap<String, Integer> countServiceMap;

    public SessionStatisticalCountWindowFunction(HashMap<String, Integer> countDstMap, HashMap<String, Integer> countServiceMap) {
        this.countDstMap = countDstMap;
        this.countServiceMap = countServiceMap;
    }

    @Override
    public void apply(TimeWindow window, Iterable<SessionResult> values, Collector<SessionResult> out) throws Exception {
        int totalDstNumber = 0, totalServiceNumber = 0;
        for (String dst : countDstMap.keySet()) {
            totalDstNumber += countDstMap.get(dst);
        }
        for (String service : countServiceMap.keySet()) {
            totalDstNumber += countServiceMap.get(service);
        }
        for (SessionResult value : values) {
            if (countDstMap.containsKey(value.getDstIP())) {
                // 一般一定在相同地址中含有
                int count = countDstMap.get(value.getDstIP());
                value.setCountWindowSameDstCount(count);
                value.setCountWindowSameDstPercent(((float) count) / ((float) totalDstNumber));
                value.setCountWindowDiffDstCount((totalDstNumber - count));
                value.setCountWindowDiffDstPercent(((float) (totalDstNumber - count)) / ((float) totalDstNumber));

            } else {
                // 如果发生了这种情况，说明计算出错了
                logger.error("目的地址map中不含有当前IP地址，请检查代码错误！！");
                value.setCountWindowSameDstCount(0);
                value.setCountWindowSameDstPercent(0);
                value.setCountWindowDiffDstCount(totalDstNumber);
                value.setCountWindowDiffDstPercent(1);
                value.setCountWindowSameServiceCount(0);
            }
            if(countServiceMap.containsKey(value.getProtocol())){
                value.setCountWindowSameServiceCount(countServiceMap.get(value.getProtocol()));
            }else {
                value.setCountWindowSameServiceCount(0);
            }
            out.collect(value);
        }

        countDstMap.clear();
        countServiceMap.clear();
    }

}